from obs import GetObjectHeader
from obs import ObsClient
import os
import traceback
import json
import random
from pathlib import Path
from urllib.parse import urlparse
from tqdm import tqdm
# requests库已不再需要，因为已强制使用OBS下载

def download_image_from_obs(obs_client, bucket_name, object_key, save_path, pbar=None):
    """
    使用OBS客户端下载图片到指定路径。

    Args:
        obs_client (ObsClient): OBS客户端实例。
        bucket_name (str): OBS Bucket的名称。
        object_key (str): OBS对象的Key（路径）。
        save_path (str): 图片在本地保存的完整路径。
        pbar (tqdm.tqdm, optional): tqdm进度条实例，用于更新状态。默认为None。

    Returns:
        bool: 如果下载成功返回True，否则返回False。
    """
    try:
        # 下载对象的附加头域，通常不需要特殊设置
        headers = GetObjectHeader()

        # 确保目标目录存在，如果不存在则创建
        Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

        # 文件下载
        resp = obs_client.getObject(bucket_name, object_key, save_path, headers=headers)

        if resp.status < 300:
            if pbar:
                pbar.set_postfix({'status': 'OBS下载成功', 'file': os.path.basename(save_path)})
            return True
        else:
            # OBS下载失败时打印详细错误信息，包括HTTP状态码
            # 确保即使 errorCode 和 errorMessage 为 None，也能打印出状态码
            status_code = resp.status if hasattr(resp, 'status') else 'N/A'
            request_id = resp.requestId if hasattr(resp, 'requestId') else 'N/A'
            error_code = resp.errorCode if hasattr(resp, 'errorCode') else 'N/A'
            error_message_detail = resp.errorMessage if hasattr(resp, 'errorMessage') else 'N/A'

            error_message = (f"OBS下载失败: HTTP Status={status_code}, requestId={request_id}, "
                             f"errorCode={error_code}, errorMessage={error_message_detail}")
            if pbar:
                # 截断错误信息以适应进度条显示
                pbar.set_postfix({'status': 'OBS下载失败', 'error': error_message[:50] + '...'})
            print(f"Error downloading {object_key}: {error_message}")
            return False
    except Exception as e:
        # 捕获下载过程中的其他异常
        error_message = f"下载过程中发生异常: {e}"
        if pbar:
            pbar.set_postfix({'status': 'OBS下载失败', 'error': error_message[:50] + '...'})
        print(f"Exception downloading {object_key}: {traceback.format_exc()}")
        return False

def process_data(input_file, output_file, image_root_dir, obs_client, default_obs_bucket_name):
    """
    处理JSON数据并下载图片，强制使用OBS进行下载。

    Args:
        input_file (str): 输入的JSON文件路径。
        output_file (str): 处理后输出的JSON文件路径。
        image_root_dir (str): 图片保存的根目录。
        obs_client (ObsClient): OBS客户端实例。
        default_obs_bucket_name (str): 如果URL中无法解析出桶名，则使用的默认OBS Bucket名称。
    """
    # 创建子目录结构 (train, val, test)
    os.makedirs(os.path.join(image_root_dir, "train"), exist_ok=True)
    os.makedirs(os.path.join(image_root_dir, "val"), exist_ok=True)
    os.makedirs(os.path.join(image_root_dir, "test"), exist_ok=True)

    print("正在读取JSON数据...")
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 打乱数据顺序，确保随机划分
    random.shuffle(data)

    # 计算训练集、验证集、测试集的划分点
    total = len(data)
    train_end = int(total * 0.8)
    val_end = train_end + int(total * 0.1)

    processed_data = []
    success_count = 0
    fail_count = 0

    print("\n开始下载图片并处理数据:")
    # 使用tqdm显示处理进度
    with tqdm(data, desc="处理进度", unit="item") as pbar:
        for i, item in enumerate(pbar):
            # 确定当前条目所属的数据集类型
            if i < train_end:
                dataset_type = "train"
            elif i < val_end:
                dataset_type = "val"
            else:
                dataset_type = "test"

            image_url = item["small_pic"]
            parsed_url = urlparse(image_url)

            # 从URL的path部分获取对象键，去除开头的斜杠和任何查询参数
            # 例如: "https://pic.justeasy.cn/202506/abc.jpg?style/waterpic" -> "202506/abc.jpg"
            object_key = parsed_url.path.lstrip('/') 
            
            # 从object_key中获取文件扩展名，如果不存在则默认为.jpg
            ext = os.path.splitext(object_key)[1] or '.jpg'
            filename = f"{item['id']}{ext}"
            relative_path = f"{dataset_type}/{filename}"  # 图片相对于image_root_dir的路径
            absolute_path = os.path.join(image_root_dir, relative_path)  # 图片在本地保存的绝对路径

            # 确定桶名：
            # 由于用户明确指出要使用OBS，并且提供了桶名，这里直接使用 default_obs_bucket_name
            current_bucket_name = default_obs_bucket_name 

            # 打印调试信息
            print(f"\n--- 正在处理URL: {image_url}")
            print(f"    本地保存路径 (Save Path): {absolute_path}")
            print(f"    尝试使用的桶名 (Bucket): {current_bucket_name}")
            print(f"    尝试使用的对象键 (Object Key): {object_key}")

            success = False
            # 强制使用OBS客户端下载
            if obs_client: # 只有当obs_client成功初始化时才尝试OBS下载
                pbar.set_postfix({'status': 'OBS下载中...', 'bucket': current_bucket_name, 'key': object_key[:30] + '...'})
                success = download_image_from_obs(obs_client, current_bucket_name, object_key, absolute_path, pbar)
            else:
                print("警告: OBS客户端未初始化，无法执行OBS下载。请检查AK/SK设置。")
                # 如果OBS客户端未初始化，这里将直接标记为失败，不会退化到HTTP下载


            if success:
                # 如果下载成功，则复制当前条目并添加新的字段
                new_item = item.copy()
                new_item["image"] = relative_path  # 存储图片的相对路径
                new_item["dataset"] = dataset_type # 存储数据集类型
                processed_data.append(new_item)
                success_count += 1
            else:
                fail_count += 1

    # 保存处理后的数据到JSON文件
    print("\n正在保存处理结果...")
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(processed_data, f, ensure_ascii=False, indent=2)

    print("\n处理完成！")
    print(f"总条目数: {len(data)}")
    print(f"成功处理: {success_count}")
    print(f"失败处理: {fail_count}")
    print(f"训练集: {train_end}, 验证集: {val_end-train_end}, 测试集: {total-val_end}")
    print(f"结果已保存到: {output_file}")
    print(f"图片目录结构:")
    print(f"  {image_root_dir}/")
    print(f"  ├── train/")
    print(f"  ├── val/")
    print(f"  └── test/")

if __name__ == "__main__":
    # 配置OBS凭据和Endpoint
    # 推荐通过环境变量获取AKSK，这里也可以使用其他外部引入方式传入。如果使用硬编码可能会存在泄露风险
    # 您可以登录访问管理管理台获取访问密钥AK/SK，获取方式请参见https://support.huaweicloud.com/usermanual-ca/ca_01_0003.html
    ak = os.getenv("AccessKeyID")  # 从环境变量获取AccessKeyID
    sk = os.getenv("SecretAccessKey") # 从环境变量获取SecretAccessKey
    # 【可选】如果使用临时AKSK和SecurityToken访问OBS，则同样推荐通过环境变量获取
    # security_token = os.getenv("SecurityToken")

    # server填写Bucket对应的Endpoint。
    # 对于南京地区，通常是 https://obs.cn-east-2.myhuaweicloud.com
    # 请根据您的实际OBS所在区域和华为云官方文档确认准确的Endpoint。
    server = "https://obs.cn-south-1.myhuaweicloud.com/" 

    # 示例Bucket名称，请根据实际情况修改
    # 如果您的JSON数据中的URL始终指向同一个OBS Bucket，可以在这里设置默认值
    # 如果URL中包含不同的Bucket，则会在process_data中尝试从URL解析，否则使用此默认值
    # 注意：当URL是自定义域名时，解析出的桶名可能不准确，此时会使用此默认桶名。
    # 请确保此默认桶名是您希望尝试下载的OBS桶。
    default_obs_bucket_name = "pic-justeasy-cn" # <-- 已根据您的输入更新为实际桶名称 !!!

    # 检查AK/SK环境变量是否已设置
    if not ak or not sk:
        print("错误: AccessKeyID 或 SecretAccessKey 环境变量未设置。")
        print("请设置环境变量 AccessKeyID 和 SecretAccessKey 后重试。")
        # exit(1) # 不再强制退出，因为可能大部分图片不是OBS下载，但仍建议设置AK/SK
        print("注意：如果您的图片URL中包含OBS标准域名，但未设置AK/SK，OBS下载将失败。")

    # 创建obsClient实例
    obsClient = None
    if ak and sk: # 只有当AK/SK都设置了才尝试创建OBS客户端
        try:
            obsClient = ObsClient(access_key_id=ak, secret_access_key=sk, server=server)
        except Exception as e:
            print(f"初始化ObsClient失败: {e}")
            print("请检查AccessKeyID, SecretAccessKey和Endpoint配置。", traceback.format_exc()) # 打印完整堆栈
            # exit(1) # 不再强制退出
            print("注意：OBS客户端初始化失败，所有OBS下载将退化为HTTP下载。")
    else:
        print("未设置AccessKeyID或SecretAccessKey，OBS客户端将不会被初始化。所有图片将尝试通过HTTP下载。")

    # 配置输入/输出文件和图片根目录路径
    input_json = "merged_output_GZ.json"  # 输入的合并JSON文件
    output_json = "final_dataset_GZ.json"  # 处理后的输出JSON文件
    image_root_directory = "dataset_images_NewGZ"  # 图片根目录（将自动创建train/val/test子目录）

    # 检查并安装tqdm库
    try:
        from tqdm import tqdm
    except ImportError:
        print("正在安装tqdm库...")
        import subprocess
        try:
            subprocess.check_call(["pip", "install", "tqdm"])
            from tqdm import tqdm
        except Exception as e:
            print(f"安装tqdm失败: {e}")
            print("请手动安装tqdm: pip install tqdm", traceback.format_exc()) # 打印完整堆栈
            exit(1) # 退出程序，因为缺少必要的库

    # 调用处理函数
    process_data(input_json, output_json, image_root_directory, obsClient, default_obs_bucket_name)

    # 关闭OBS客户端，释放资源
    if obsClient:
        obsClient.close()
